{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can directly use model client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "api_kwargs: {'model': 'gpt-3.5-turbo', 'temperature': 0.5, 'max_tokens': 100, 'messages': [{'role': 'system', 'content': 'User: What is the capital of France?\\n'}]}\n",
      "response_text: The capital of France is Paris.\n",
      "api_kwargs: {'model': 'text-embedding-3-small', 'dimensions': 8, 'encoding_format': 'float', 'input': ['What is the capital of France?', 'What is the capital of France?']}\n",
      "reponse_embedder_output: EmbedderOutput(data=[Embedding(embedding=[0.6175549, 0.24047995, 0.4509756, 0.37041178, -0.33437008, -0.050995983, -0.24366009, 0.21549304], index=0), Embedding(embedding=[0.6175549, 0.24047995, 0.4509756, 0.37041178, -0.33437008, -0.050995983, -0.24366009, 0.21549304], index=1)], model='text-embedding-3-small', usage=Usage(prompt_tokens=14, total_tokens=14), error=None, raw_response=None)\n"
     ]
    }
   ],
   "source": [
    "from lightrag.components.model_client import OpenAIClient\n",
    "from lightrag.core.types import ModelType\n",
    "from lightrag.utils import setup_env\n",
    "\n",
    "openai_client = OpenAIClient()\n",
    "\n",
    "query = \"What is the capital of France?\"\n",
    "\n",
    "# try LLM model\n",
    "model_type = ModelType.LLM\n",
    "\n",
    "prompt = f\"User: {query}\\n\"\n",
    "model_kwargs = {\"model\": \"gpt-3.5-turbo\", \"temperature\": 0.5, \"max_tokens\": 100}\n",
    "api_kwargs = openai_client.convert_inputs_to_api_kwargs(input=prompt, \n",
    "                                                        model_kwargs=model_kwargs, \n",
    "                                                        model_type=model_type)\n",
    "print(f\"api_kwargs: {api_kwargs}\")\n",
    "\n",
    "response = openai_client.call(api_kwargs=api_kwargs, model_type=model_type)\n",
    "response_text = openai_client.parse_chat_completion(response)\n",
    "print(f\"response_text: {response_text}\")\n",
    "\n",
    "# try embedding model\n",
    "model_type = ModelType.EMBEDDER\n",
    "# do batch embedding\n",
    "input = [query] * 2\n",
    "model_kwargs = {\"model\": \"text-embedding-3-small\", \"dimensions\": 8, \"encoding_format\": \"float\"}\n",
    "api_kwargs = openai_client.convert_inputs_to_api_kwargs(input=input, model_kwargs=model_kwargs, model_type=model_type)\n",
    "print(f\"api_kwargs: {api_kwargs}\")\n",
    "\n",
    "\n",
    "\n",
    "response = openai_client.call(api_kwargs=api_kwargs, model_type=model_type)\n",
    "reponse_embedder_output = openai_client.parse_embedding_response(response)\n",
    "print(f\"reponse_embedder_output: {reponse_embedder_output}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "my-project-kernel",
   "language": "python",
   "name": "my-project-kernel"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
